package com.wt.admin.service.ai.impl.agents.easyai.api;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.tool.ToolCallingChatOptions;

import java.util.*;
import java.util.stream.Collectors;

public class EasyAIOptions implements ToolCallingChatOptions {

    @JsonIgnore
    private Boolean internalToolExecutionEnabled;
    @JsonIgnore
    private List<FunctionCallback> toolCallbacks = new ArrayList();
    @JsonIgnore
    private Set<String> toolNames = new HashSet();
    @JsonIgnore
    private Map<String, Object> toolContext = new HashMap();
    private static final List<String> NON_SUPPORTED_FIELDS = List.of("model", "format", "keep_alive", "truncate");
    @JsonProperty("format")
    private Object format;
    @JsonProperty("model")
    private String model;
    @JsonProperty("frequency_penalty")
    private Double frequencyPenalty;
    @JsonProperty("num_predict")
    private Integer maxTokens;
    @JsonProperty("presence_penalty")
    private Double presencePenalty;
    @JsonProperty("stop")
    private List<String> stopSequences;
    @JsonProperty("temperature")
    private Double temperature;
    @JsonProperty("top_k")
    private Integer topK;
    @JsonProperty("top_p")
    private Double topP;
    @JsonProperty("keep_alive")
    private String keepAlive;

    @Override
    public String getModel() {
        return this.model;
    }

    public void setModel(String model) {
        this.model = model;
    }
    @Override
    public Double getFrequencyPenalty() {
        return this.frequencyPenalty;
    }

    public void setFrequencyPenalty(Double frequencyPenalty) {
        this.frequencyPenalty = frequencyPenalty;
    }
    @Override
    public Integer getMaxTokens() {
        return this.maxTokens;
    }

    public void setMaxTokens(Integer maxTokens) {
        this.maxTokens = maxTokens;
    }
    @Override
    public Double getPresencePenalty() {
        return this.presencePenalty;
    }

    public void setPresencePenalty(Double presencePenalty) {
        this.presencePenalty = presencePenalty;
    }
    @Override
    public List<String> getStopSequences() {
        return this.stopSequences != null ? Collections.unmodifiableList(this.stopSequences) : null;
    }

    public void setStopSequences(List<String> stopSequences) {
        this.stopSequences = stopSequences;
    }
    @Override
    public Double getTemperature() {
        return this.temperature;
    }

    public void setTemperature(Double temperature) {
        this.temperature = temperature;
    }
    @Override
    public Integer getTopK() {
        return this.topK;
    }

    public void setTopK(Integer topK) {
        this.topK = topK;
    }
    @Override
    public Double getTopP() {
        return this.topP;
    }

    public void setTopP(Double topP) {
        this.topP = topP;
    }

    public Object getFormat() {
        return this.format;
    }
    public void setFormat(Object format) {
        this.format = format;
    }
    public String getKeepAlive() {
        return this.keepAlive;
    }

    public void setKeepAlive(String keepAlive) {
        this.keepAlive = keepAlive;
    }

    public static Map<String, Object> filterNonSupportedFields(Map<String, Object> options) {
        return options.entrySet().stream()
                .filter((e) -> !NON_SUPPORTED_FIELDS.contains(e.getKey()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }
    public Map<String, Object> toMap() {
        return ModelOptionsUtils.objectToMap(this);
    }

    @Override
    public <T extends ChatOptions> T copy() {
        EasyAIOptions copy = new EasyAIOptions();
        copy.setModel(this.getModel());
        copy.setFrequencyPenalty(this.getFrequencyPenalty());
        copy.setMaxTokens(this.getMaxTokens());
        copy.setPresencePenalty(this.getPresencePenalty());
        copy.setStopSequences(this.getStopSequences() != null ? new ArrayList(this.getStopSequences()) : null);
        copy.setTemperature(this.getTemperature());
        copy.setTopK(this.getTopK());
        copy.setTopP(this.getTopP());
        return (T) copy;
    }



    @JsonIgnore
    @Override
    public List<FunctionCallback> getToolCallbacks() {
        return this.toolCallbacks;
    }

    @JsonIgnore
    @Override
    public void setToolCallbacks(List<FunctionCallback> toolCallbacks) {
        this.toolCallbacks = toolCallbacks;
    }

    @JsonIgnore
    @Override
    public Set<String> getToolNames() {
        return this.toolNames;
    }

    @JsonIgnore
    @Override
    public void setToolNames(Set<String> toolNames) {
        this.toolNames = toolNames;
    }

    @JsonIgnore
    @Override
    public Boolean isInternalToolExecutionEnabled() {
        return this.internalToolExecutionEnabled;
    }

    @JsonIgnore
    @Override
    public void setInternalToolExecutionEnabled(Boolean internalToolExecutionEnabled) {
        this.internalToolExecutionEnabled = internalToolExecutionEnabled;
    }

    @Override
    @JsonIgnore
    public List<FunctionCallback> getFunctionCallbacks() {
        return this.getToolCallbacks();
    }

    @Override
    @JsonIgnore
    public void setFunctionCallbacks(List<FunctionCallback> functionCallbacks) {
        this.setToolCallbacks(functionCallbacks);
    }

    @Override
    @JsonIgnore
    public Set<String> getFunctions() {
        return this.getToolNames();
    }

    @Override
    @JsonIgnore
    public void setFunctions(Set<String> functions) {
        this.setToolNames(functions);
    }
    @Override
    @JsonIgnore
    public Map<String, Object> getToolContext() {
        return this.toolContext;
    }
    @Override
    @JsonIgnore
    public void setToolContext(Map<String, Object> toolContext) {
        this.toolContext = toolContext;
    }
}
